package com.lcj.ai.controller;


import com.lcj.ai.config.DogScheduler;
import com.lcj.ai.entity.Dog;
import com.lcj.ai.mapper.DogMapper;
import com.lcj.ai.mapper.VectorStoreMapper;
import io.modelcontextprotocol.client.McpSyncClient;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.client.advisor.PromptChatMemoryAdvisor;
import org.springframework.ai.chat.client.advisor.vectorstore.QuestionAnswerAdvisor;
import org.springframework.ai.chat.memory.ChatMemory;
import org.springframework.ai.chat.memory.MessageWindowChatMemory;
import org.springframework.ai.chat.memory.repository.jdbc.JdbcChatMemoryRepository;
import org.springframework.ai.document.Document;
import org.springframework.ai.mcp.SyncMcpToolCallbackProvider;
import org.springframework.ai.vectorstore.VectorStore;
import org.springframework.web.bind.annotation.*;

import java.util.List;
import java.util.Map;

@RestController
@RequestMapping("/ai")
public class AiController {
    private final ChatClient chatClient;

    private final DogMapper dogMapper;

    public AiController(ChatClient.Builder builder, JdbcChatMemoryRepository jdbcChatMemoryRepository,
                        VectorStore vectorStore, DogMapper dogMapper, VectorStoreMapper vectorStoreMapper,
                        /*DogScheduler dogScheduler*/ McpSyncClient mcpSyncClient, DogMapper dogMapper1) {
        this.dogMapper = dogMapper1;
        // 检查向量数据库是否已经有数据，避免重复写入
        var count = vectorStoreMapper.selectCount(null);
        if (count == 0) {
            dogMapper.selectList(null).forEach(dog -> {
                // 构建 Document，内容可以自定义，这里包含了ID、名字和描述
                var dogDocument = new Document("id: %s, name: %s, description: %s".formatted(
                        dog.getId(), dog.getName(), dog.getDescription()
                ), Map.of("id", dog.getId().toString()));
                vectorStore.add(List.of(dogDocument)); // 添加到向量数据库
            });
        }

        var systemPromptText = """
                你是一个 AI 助手，帮助人们从名为“Pooch Palace”的狗狗领养机构领养狗狗。
                该机构在以下城市设有分部：里约热内卢、墨西哥城、首尔、东京、新加坡、纽约、阿姆斯特丹、巴黎、孟买、新德里、巴塞罗那、伦敦和旧金山。
                可领养狗狗的信息会展示在下方。如果没有信息，请礼貌地回复说我们目前没有可领养的狗狗。
                """;
        this.chatClient = builder.
                defaultSystem(systemPromptText).
                defaultAdvisors(
                        PromptChatMemoryAdvisor.builder(
                                MessageWindowChatMemory.builder().chatMemoryRepository(jdbcChatMemoryRepository).build()
                        ).build(), // 默认的聊天记忆顾问
                        new QuestionAnswerAdvisor(vectorStore) // 问答顾问
                )/*.defaultTools(dogScheduler)*/ // 默认的工具
                .defaultToolCallbacks(new SyncMcpToolCallbackProvider(mcpSyncClient))
                .build();
    }

    @GetMapping("/{user}/chat")
    public String chat(@PathVariable("user") String user, @RequestParam("message") String message) {
        return chatClient.
                prompt().
                user(message).
                advisors(a -> a.param(ChatMemory.CONVERSATION_ID, user)).
                call().
                content();
    }

    @GetMapping("/testData")
    public String testData() {
        // 插入测试数据
        dogMapper.insert(new Dog("旺财", "张三", "一只可爱的小土狗"));
        dogMapper.insert(new Dog("大黄", "李四", "一只聪明的rottweiler"));
        dogMapper.insert(new Dog("小黑", "王五", "一只friendlyrottweiler"));
        dogMapper.insert(new Dog("小黄", "赵六", "一只friendlyrottweiler"));
        dogMapper.insert(new Dog("小红", "孙七", "一只friendlyrottweiler"));
        dogMapper.insert(new Dog("小蓝", "周八", "一只friendlyrottweiler"));
        return "插入测试数据";
    }
}
